"""
Case Type   : DATAVEC模块
Case Name   : 存在HNSWPQ l2距离索引的情况下，关闭enable_pq，验证表数据插入及查询
Create At   : 2025/4/10
Owner       : @peilinqian
Description :
    1. 修改参数enable_pq为true;
    2. 重启数据库，使参数生效;
    3. 创建表及pq索引并插入数据;
    4. 主机使用pq索引查询数据;
    5. 备机使用pq索引查询数据;
    6. 修改参数enable_pq为off;
    7. 重启数据库，使参数生效;
    8. 主机使用pq索引查询数据;
    9. 备机使用pq索引查询数据;
    10. 主机插入数据;
Expect      :
    1. 修改参数enable_pq为true; 成功
    2. 重启数据库，使参数生效; 重启成功，参数生效;
    3. 创建表及pq索引并插入数据; 成功
    4. 主机使用pq索引查询数据; 成功
    5. 备机使用pq索引查询数据; 成功
    6. 修改参数enable_pq为off; 成功
    7. 重启数据库，使参数生效; 重启成功，参数生效;
    8. 主机使用pq索引查询数据; 失败
    9. 备机使用pq索引查询数据; 失败
    10. 主机插入数据; 失败
History     :
"""
import os
import unittest

from testcase.utils.Common import Common
from testcase.utils.CommonSH import CommonSH
from testcase.utils.Constant import Constant
from testcase.utils.Logger import Logger
from yat.test import Node
from yat.test import macro

pri_dbuser = Node(node='PrimaryDbUser')
pq_lib_path = '/usr/local/sra_recall/lib/'
log = Logger()
log.info("----判断本地环境是否存在pq插件----")
check_cmd = f'if [ -d {pq_lib_path} ]; ' \
            f'then echo "True"; else echo "False"; fi'
log.info(check_cmd)
check_result = pri_dbuser.sh(check_cmd).result()
log.info(check_result)


@unittest.skipIf(check_result == 'False',
                 '----环境不具备PQ插件，跳过执行用例----')
class DataVecTestCase(unittest.TestCase):
    def setUp(self):
        log.info(f'----{os.path.basename(__file__)}:start----')
        self.sh = CommonSH('PrimaryDbUser')
        self.sta_1_sh = CommonSH('Standby1DbUser')
        self.sta_2_sh = CommonSH('Standby2DbUser')
        self.constant = Constant()
        self.com = Common()
        self.tb_name = 't_datavec_pq_case0004'
        self.pq_env_name = 'DATAVEC_PQ_LIB_PATH'
        step_txt = '----step0:设置pq环境变量; expect:成功----'
        log.info(step_txt)
        check_env_cmd = f'if grep -q "{self.pq_env_name}" ' \
                        f'"{macro.DB_ENV_PATH}"; ' \
                        f'then echo "True" ; else echo "False" ;fi'
        set_env_cmd = f"echo 'export {self.pq_env_name}={pq_lib_path}' " \
                      f">> {macro.DB_ENV_PATH}"
        for node in [self.sh.node, self.sta_1_sh.node, self.sta_2_sh.node]:
            log.info(check_env_cmd)
            check_env_result = node.sh(check_env_cmd).result()
            log.info(check_env_result)
            if check_env_result == 'False':
                log.info('--设置pq环境变量--')
                log.info(set_env_cmd)
                try:
                    result = node.sh(set_env_cmd).result()
                    log.info(f"{node.name} 设置结果: {result}")
                except Exception as e:
                    log.error(f"{node.name} 设置异常: {str(e)}")
            check_env_result = node.sh(check_env_cmd).result()
            log.info(f"{node.name} 环境变量检查结果: {check_env_result}")
            self.assertEqual(check_env_result, 'True', '执行失败:' + step_txt)

        step_txt = '----step0:查看enable_pq初始配置值;----'
        log.info(step_txt)
        self.init_para = self.com.show_param("enable_pq")

    def test_main(self):
        step_txt = '----step1:修改参数enable_pq为true; expect:成功----'
        log.info(step_txt)
        msg = self.sh.execute_gsguc('set',
                                    self.constant.GSGUC_SUCCESS_MSG,
                                    f"enable_pq= 1")
        self.assertTrue(msg, '执行失败:' + step_txt)

        step_txt = '----step2:重启数据库，使参数生效; expect:重启成功，参数生效;----'
        log.info(step_txt)
        restart_result = self.sh.restart_db_cluster()
        self.assertTrue(restart_result)
        log.info('----查询参数----')
        self.new_para = self.com.show_param("enable_pq")
        self.assertEqual(self.new_para, 'on', '执行失败:' + step_txt)

        step_txt = '----step3:创建表及pq索引并插入数据; expect:成功----'
        log.info(step_txt)
        create_tb = f"drop table if exists {self.tb_name};" \
                    f"create table {self.tb_name} (id bigserial primary key," \
                    f"val1 vector(3));"
        insert_sql = f"insert into {self.tb_name}(val1) values" \
                     f"(''['||(random()*1000+1)::int||'," \
                     f"'||(random()*1000+1)::int||'," \
                     f"'||(random()*1000+1)::int||']'')"
        insert_block = f''' declare
        v_sql varchar:='';
        begin for i in 1..10000 loop
        v_sql ='{insert_sql}';
        execute v_sql;
        end loop;
        end;'''
        create_index = f"create index on {self.tb_name} using hnsw " \
                       f"(val1 vector_l2_ops) with (enable_pq=on, pq_m=1);"
        select_sql = f"set enable_seqscan=off;" \
                     f"explain select val1,val1<->'[3,3,3]' from " \
                     f"{self.tb_name} order by val1<-> '[3,3,3]' limit 10;" \
                     f"select val1,val1<->'[3,3,3]' from " \
                     f"{self.tb_name} order by val1<-> '[3,3,3]' limit 10;"
        sql = f'''
        --创建表
        {create_tb}
        --插入数据量较大
        {insert_block}
        --创建索引
        {create_index}
        '''
        result = self.sh.execut_db_sql(sql)
        log.info(result)
        self.assertIn(self.constant.TABLE_CREATE_SUCCESS, result,
                      '执行失败:' + step_txt)

        step_txt = '----step4:主机使用pq索引查询数据; expect:成功----'
        log.info(step_txt)
        result = self.sta_1_sh.execut_db_sql(select_sql)
        log.info(result)
        self.assertIn("Ann Index Scan using", result, '执行失败:' + step_txt)
        self.assertIn("10 rows", result, '执行失败:' + step_txt)

        step_txt = '----step5:备机使用pq索引查询数据; expect:成功----'
        log.info(step_txt)
        result = self.sta_1_sh.execut_db_sql(select_sql)
        log.info(result)
        self.assertIn("Ann Index Scan using", result, '执行失败:' + step_txt)
        self.assertIn("10 rows", result, '执行失败:' + step_txt)

        step_txt = '----step6:修改参数enable_pq为off; expect:成功----'
        log.info(step_txt)
        msg = self.sh.execute_gsguc('set',
                                    self.constant.GSGUC_SUCCESS_MSG,
                                    f"enable_pq= 0")
        self.assertTrue(msg, '执行失败:' + step_txt)

        step_txt = '----step7:重启数据库，使参数生效; expect:重启成功，参数生效;----'
        log.info(step_txt)
        restart_result = self.sh.restart_db_cluster()
        self.assertTrue(restart_result)
        log.info('----查询参数----')
        self.new_para = self.com.show_param("enable_pq")
        self.assertEqual(self.new_para, 'off', '执行失败:' + step_txt)

        step_txt = '----step8:主机使用pq索引查询数据; expect:失败----'
        log.info(step_txt)
        result = self.sh.execut_db_sql(select_sql)
        log.info(result)
        self.assertIn("ERROR:  the SQL involves operations related to HNSWPQ",
                      result, '执行失败:' + step_txt)

        step_txt = '----step9:备机使用pq索引查询数据; expect:失败----'
        log.info(step_txt)
        result = self.sta_1_sh.execut_db_sql(select_sql)
        log.info(result)
        self.assertIn("ERROR:  the SQL involves operations related to HNSWPQ",
                      result, '执行失败:' + step_txt)

        step_txt = '----step10:主机插入数据; expect:失败----'
        log.info(step_txt)
        result = self.sh.execut_db_sql(insert_block)
        log.info(result)
        self.assertIn("ERROR:  the SQL involves operations related to HNSWPQ",
                      result, '执行失败:' + step_txt)

    def tearDown(self):
        step_txt = '----this is teardown----'
        log.info(step_txt)
        step1_txt = '----teardown:清理测试表; expect:成功----'
        log.info(step1_txt)
        drop_cmd = f'drop table if exists {self.tb_name}'
        drop_result = self.sh.execut_db_sql(drop_cmd)

        step2_txt = '----teardown:还原参数enable_pq为初始值; ' \
                    'expect:修改成功----'
        log.info(step2_txt)
        msg = self.sh.execute_gsguc('set',
                                    self.constant.GSGUC_SUCCESS_MSG,
                                    f"enable_pq={self.init_para}")
        step3_txt = '----teardown:重启数据库，使参数生效; expect:重启成功----'
        log.info(step3_txt)
        restart_result = self.sh.restart_db_cluster()

        log.info(f'----{os.path.basename(__file__)}:end----')
        self.assertIn(self.constant.TABLE_DROP_SUCCESS, drop_result,
                      '执行失败:' + step_txt + step1_txt)
        self.assertTrue(msg, '执行失败:' + step_txt + step2_txt)
        self.assertTrue(restart_result, '执行失败:' + step_txt + step3_txt)
